import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from torchvision.transforms import Compose,ToTensor,Normalize
from torchvision import datasets
import os

'''
    ! Attentions of coding GAN:
    1. detach的使用,错误使用会导致gen_img从计算图中分离出来.
    2. model不可以作为class的属性,会导致没办法进行优化.
'''

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        def block(in_features, out_features, normalize=True):
            layers = [nn.Linear(in_features, out_features)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_features,0.8))
            layers.append(nn.LeakyReLU(0,2))
            return layers
        self.model = nn.Sequential(
            *block(100, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Linear(1024, 28*28*1),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.model(z)
        img = img.view(-1, 1, 28, 28)
        return img


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator,self).__init__()
        self.model = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        img_flat = x.reshape(-1, 28*28)
        prob = self.model(img_flat)
        return prob

class OriginGAN():
    def __init__(self):
        self.device = "cuda:0"

    def train(self):
        dataloader = DataLoader(
            datasets.MNIST(
                "./data/mnist",
                train=True,
                download=True,
                transform=Compose([ToTensor(), Normalize(mean=0.5, std=0.5)])
            ),
            batch_size=256,
            shuffle=True
        )

        generator = Generator().to(self.device)
        discriminator = Discriminator().to(self.device)
        optimizer_G = torch.optim.Adam(generator.parameters(),lr=0.0002)
        optimizer_D = torch.optim.Adam(discriminator.parameters(),lr=0.0002)

        adversarial_loss = nn.BCELoss().to(self.device)
        if os.path.exists("./model/generator.pth") and os.path.exists("./model/discriminator.pth"):
            generator.load_state_dict(torch.load("./model/generator.pth"))
            discriminator.load_state_dict(torch.load("./model/discriminator.pth"))
            print("model loaded.")
        for epoch in range(1000):
            epoch_g_loss =0.
            epoch_d_loss = 0.
            for i, (real_imgs, _) in enumerate(dataloader):
                valid = torch.ones((real_imgs.shape[0], 1),dtype=torch.float32, requires_grad=False).to(self.device)
                fake = torch.zeros((real_imgs.shape[0], 1),dtype=torch.float32, requires_grad=False).to(self.device)
                real_imgs = real_imgs.to(self.device)

                # -------------- # 
                # train generator
                # -------------- #
                optimizer_G.zero_grad()
                z = torch.randn((real_imgs.shape[0], 100)).to(self.device)
                gen_imgs = generator(z)
                g_loss = adversarial_loss(discriminator(gen_imgs), valid)
                g_loss.backward()
                optimizer_G.step()
                epoch_g_loss += g_loss.item()

                # ------------------- #
                # train discriminator
                # ------------------- #
                optimizer_D.zero_grad()

                real_loss = adversarial_loss(discriminator(real_imgs), valid)
                fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)
                d_loss = (real_loss + fake_loss)/2

                d_loss.backward()
                optimizer_D.step()
                epoch_d_loss += d_loss.item()
            print("epoch: ",epoch,"\tg_loss: ",epoch_g_loss,"\td_loss:",epoch_d_loss)
            if epoch % 10 == 0:
                save_image(gen_imgs[:25], 'res/'+str(epoch)+'.png', nrow=5, normalize=True)
                torch.save(generator.state_dict(), "./model/generator.pth")
                torch.save(discriminator.state_dict(), "./model/discriminator.pth")
if __name__ == "__main__":
    gan = OriginGAN()
    gan.train()

                

        
        